# 参数配置
import argparse   # argparse：为py文件封装好可以选择的参数，使他们更加灵活，丰富


def parse_args():
    parser = argparse.ArgumentParser(description='Intelligence Tang Poetry Generator')
    help_ = 'You can set this value in terminal --write value can be poem or lyric'
    parser.add_argument('-w', '--write', default='poem', choices=['poem', 'lyric'],help=help_)
    help_ = 'Choose to train or generate'
    parser.add_argument('--train', dest='train', action='store_true', help=help_)
    parser.add_argument('--no-train', dest='train', action='store_false', help=help_)
    parser.set_defaults(train=False)
    args_ = parser.parse_args()
    return args_


if __name__ == '__main__':
    # print('ccc')

    args = parse_args()    # 解析命令行传来的参数
    print(args)

    if args.write == 'poem':
        from inference import TangPoems
        if args.train:
            TangPoems.main(True)
        else:
            # print("主函数111", args)
            TangPoems.main(False)
            # print("主函数", args)
    elif args.write == 'lyric':
        from inference import SongLyrics
        print(args.train)
        if args.train:
            SongLyrics.main(True)
        else:
            SongLyrics.main(False)
    else:
        print("[INFO] Write option can only be TangPoetry or SongLyrics")

